#!/usr/bin/env python3
"""
Add team_id column to metrics table and initialize with data from metric_team_assignments
"""

import sqlite3
import os
from typing import Dict, Optional

def get_database_path() -> str:
    """Get the path to the SQLite database"""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    backend_data_dir = os.path.join(current_dir, 'data')
    db_path = os.path.join(backend_data_dir, 'uma_insight.db')
    if not os.path.exists(db_path):
        project_root = os.path.dirname(current_dir)
        db_path = os.path.join(project_root, 'backend', 'data', 'uma_insight.db')
    return db_path

def migrate_metrics_table():
    """Add team_id column to metrics table"""
    db_path = get_database_path()

    if not os.path.exists(db_path):
        print(f"❌ Database not found at: {db_path}")
        return False

    print(f"🔄 [Migration] Starting migration for metrics table...")
    print(f"📁 Database path: {db_path}")

    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # Check if team_id column already exists
        cursor.execute("PRAGMA table_info(metrics)")
        columns = [column[1] for column in cursor.fetchall()]

        if 'team_id' in columns:
            print("✅ [Migration] team_id column already exists in metrics table")
            return True

        # Add team_id column as foreign key to teams table
        print("🔧 [Migration] Adding team_id column to metrics table...")
        cursor.execute("""
            ALTER TABLE metrics
            ADD COLUMN team_id TEXT
        """)

        # Create index for team_id
        print("🔧 [Migration] Creating index for team_id...")
        cursor.execute("""
            CREATE INDEX idx_metrics_team_id ON metrics(team_id)
        """)

        # Add foreign key constraint (SQLite doesn't support adding FK constraints to existing tables easily,
        # so we'll rely on application-level validation)
        print("🔧 [Migration] team_id column added successfully")

        conn.commit()
        print("✅ [Migration] Metrics table schema updated successfully")
        return True

    except sqlite3.Error as e:
        print(f"❌ [Migration] Database error: {e}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

def initialize_team_id_data():
    """Initialize team_id data from metric_team_assignments table"""
    db_path = get_database_path()

    print("🔄 [Initialization] Starting team_id data initialization...")

    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # Check if we have data to migrate
        cursor.execute("SELECT COUNT(*) FROM metric_team_assignments")
        assignment_count = cursor.fetchone()[0]

        if assignment_count == 0:
            print("✅ [Initialization] No metric_team_assignments data found, skipping initialization")
            return True

        print(f"📊 [Initialization] Found {assignment_count} team assignments to process")

        # Get team name to ID mapping from teams table
        cursor.execute("SELECT id, name FROM teams")
        team_mapping = {name: team_id for team_id, name in cursor.fetchall()}
        print(f"🏢 [Initialization] Found {len(team_mapping)} teams in database")

        if not team_mapping:
            print("⚠️ [Initialization] No teams found in teams table")
            return True

        # Process metric_team_assignments to determine team_id for each metric
        # Strategy: Use the most recent assignment for each metric
        cursor.execute("""
            SELECT metric_id, team_name,
                   ROW_NUMBER() OVER (PARTITION BY metric_id ORDER BY assigned_at DESC) as rn
            FROM metric_team_assignments
        """)

        assignments = cursor.fetchall()
        metrics_team_updates = []

        for metric_id, team_name, row_num in assignments:
            if row_num == 1 and team_name in team_mapping:  # Only process the most recent assignment
                team_id = team_mapping[team_name]
                metrics_team_updates.append((team_id, metric_id))

        print(f"📝 [Initialization] Processing {len(metrics_team_updates)} unique metrics")

        # Update metrics table with team_id
        updated_count = 0
        for team_id, metric_id in metrics_team_updates:
            cursor.execute("""
                UPDATE metrics
                SET team_id = ?
                WHERE id = ?
            """, (team_id, metric_id))

            if cursor.rowcount > 0:
                updated_count += 1

        conn.commit()

        print(f"✅ [Initialization] Updated team_id for {updated_count} metrics")
        print(f"📊 [Initialization] Team distribution:")

        # Show team distribution
        cursor.execute("""
            SELECT t.name, COUNT(m.id) as metric_count
            FROM metrics m
            JOIN teams t ON m.team_id = t.id
            WHERE m.team_id IS NOT NULL
            GROUP BY t.id, t.name
            ORDER BY metric_count DESC
            LIMIT 10
        """)

        for row in cursor.fetchall():
            team_name, count = row
            print(f"   {team_name}: {count} metrics")

        return True

    except sqlite3.Error as e:
        print(f"❌ [Initialization] Database error: {e}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

def verify_migration():
    """Verify the migration was successful"""
    db_path = get_database_path()

    print("🔍 [Verification] Verifying migration...")

    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # Check metrics table structure
        cursor.execute("PRAGMA table_info(metrics)")
        columns = [column[1] for column in cursor.fetchall()]

        if 'team_id' not in columns:
            print("❌ [Verification] team_id column not found in metrics table")
            return False

        print("✅ [Verification] team_id column found in metrics table")

        # Check team_id data
        cursor.execute("SELECT COUNT(*) FROM metrics WHERE team_id IS NOT NULL")
        team_assigned_count = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM metrics")
        total_metrics_count = cursor.fetchone()[0]

        print(f"📊 [Verification] Metrics with team_id: {team_assigned_count}/{total_metrics_count}")
        print(f"📊 [Verification] Team assignment coverage: {(team_assigned_count/total_metrics_count)*100:.1f}%")

        if team_assigned_count > 0:
            print("✅ [Verification] team_id data successfully initialized")
        else:
            print("⚠️ [Verification] No team_id data found (this may be expected)")

        return True

    except sqlite3.Error as e:
        print(f"❌ [Verification] Database error: {e}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

def main():
    """Main migration function"""
    print("🚀 [Team ID Migration] Starting metrics table optimization")
    print("=" * 60)

    # Step 1: Add team_id column
    if not migrate_metrics_table():
        print("❌ [Migration] Failed to add team_id column")
        return False

    print()

    # Step 2: Initialize team_id data
    if not initialize_team_id_data():
        print("❌ [Migration] Failed to initialize team_id data")
        return False

    print()

    # Step 3: Verify migration
    if not verify_migration():
        print("❌ [Migration] Migration verification failed")
        return False

    print("=" * 60)
    print("🎉 [Team ID Migration] Migration completed successfully!")
    print("💡 [Next Steps] Update API logic to use metrics.team_id instead of joining metric_team_assignments")

    return True

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)